from enum import Enum
from typing import Optional, List, Dict, Any
from pydantic import BaseModel




class ObjectiveType(str, Enum):
    PRIVACY = "PRIVACY"          # 隐私保护程度
    UTILITY = "UTILITY"          # 数据效用
    FAIRNESS = "FAIRNESS"       # 公平性
    EFFICIENCY = "EFFICIENCY"    # 效率

class DataPartyInfo(BaseModel):
    party_id: str               # 参与方ID
    data_size: int             # 数据量大小
    data_types: List[str]      # 数据类型列表
    privacy_budget: float      # 隐私预算
    min_share_ratio: float     # 最小共享比例

class MedicalShareInputParams(BaseModel):
    data_parties: List[DataPartyInfo]      # 数据参与方信息列表
    objective_weights: Dict[ObjectiveType, float]  # 目标权重字典
    max_iterations: int = 100               # 最大迭代次数
    convergence_threshold: float = 0.001    # 收敛阈值
    privacy_threshold: float = 0.8          # 隐私保护阈值

class SharingStrategy(BaseModel):
    party_id: str               # 参与方ID
    share_ratio: float         # 共享比例
    selected_features: List[str]  # 选定的特征列表
    privacy_score: float       # 隐私得分

class TimeStepResult(BaseModel):
    iteration: int             # 迭代次数
    strategies: List[SharingStrategy]  # 共享策略列表
    objective_scores: Dict[ObjectiveType, float]  # 各目标的得分
    pareto_optimal: bool       # 是否为帕累托最优解

class MedicalShareOutputParams(BaseModel):
    results: List[TimeStepResult]          # 每个时间步的结果列表
    algorithm: str = "MedicalShare"        # 算法名称
    final_strategies: List[SharingStrategy]  # 最终的共享策略列表
    convergence_history: List[float]       # 收敛历史记录
    parameters: Dict[str, Any]             # 算法参数字典

class InputParams(BaseModel):
    medical_share_params: Optional[MedicalShareInputParams] = None  # 医疗共享输入参数

class OutputParams(BaseModel):
    medical_share_results: Optional[MedicalShareOutputParams] = None  # 医疗共享输出结果 
# 1. 任务状态
class TaskStatus(str, Enum):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"
# 5. 算法请求
class AlgorithmRequest(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    input_params: InputParams


# 6. 算法中间响应
class AlgorithmMiddleResponse(BaseModel):
    task_id: str
    task_callback_url: str
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None
    input_params: InputParams
    error_message: Optional[str] = None


class AlgorithmResponse(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None

    error_message: Optional[str] = None

    output_params: OutputParams

